from torchvision.models import resnet34
import torch.nn as nn
import torch

# 
# Preconv1 BN relu maxpool 【浅层编码】 =》64
# Layer1，Layer2，Layer3，layer4
# === Backbone
# neck 
# Head 

class R34(nn.Module):
    def __init__(self, n_class):
        super().__init__()
        self.model = resnet34(pretrained=True) # Imageset 
        #conv1 = model.conv1
        self.conv1 = nn.Conv2d(in_channels=1,out_channels=64,kernel_size=3)
        self.bn1 = self.model.bn1
        self.relu = self.model.relu
        self.maxpool = self.model.maxpool
        self.avgpool = self.model.avgpool
        self.fc = nn.Linear(512, n_class) # 
        res1 = self.model.layer1
        res2 = self.model.layer2
        res3 = self.model.layer3
        res4 = self.model.layer4
        self.layers = [res1,res2,res3,res4]

        for m in self.modules():
            if isinstance(m, nn.BatchNorm2d):
                for p in m.parameters():
                    p.requires_grad = False 

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.maxpool(x)
        for l in self.layers:
            x = l(x)
        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        pred = self.fc(x)
        return pred


model = R34(10)
model.cuda()
out = model(torch.randn(1,1,28,28).cuda())
print(out.shape)
